const fs = require('fs').promises;
const path = require('path');
const { existsSync, lstatSync } = require('fs');

// 配置：需要统计的文件类型
const targetExtensions = ['.js', '.json', '.ux', '.css', '.less', '.html', '.ts'];
// 体积格式化（字节转KB）
const formatSize = (bytes) => (bytes / 1024).toFixed(2);
// 百分比格式化
const formatPercent = (percent) => percent.toFixed(1);


/**
 * 递归扫描目录，收集目标文件
 */
async function scanDirectory(dir) {
    const files = [];
    try {
        if (!existsSync(dir) || !lstatSync(dir).isDirectory()) {
            console.error(`\x1B[31m错误：${dir} 不是有效目录\x1B[0m`);
            return files;
        }

        const entries = await fs.readdir(dir, { withFileTypes: true }) || [];
        for (const entry of entries) {
            const fullPath = path.join(dir, entry.name);
            
            if (entry.name.startsWith('.')) continue;

            if (entry.isDirectory()) {
                const subFiles = await scanDirectory(fullPath);
                files.push(...subFiles);
            } else if (entry.isFile()) {
                const ext = path.extname(entry.name).toLowerCase();
                if (targetExtensions.includes(ext)) {
                    try {
                        const stat = await fs.stat(fullPath);
                        files.push({
                            dir: dir,
                            size: stat.size
                        });
                    } catch (err) {
                        console.error(`\x1B[31m无法访问文件 ${entry.name}：${err.message}\x1B[0m`);
                    }
                }
            }
        }
    } catch (err) {
        console.error(`\x1B[31m扫描目录 ${dir} 失败：${err.message}\x1B[0m`);
    }
    return files;
}


/**
 * 按目录分组统计体积
 */
function calculateDirSizes(files, rootDir) {
    const dirSizes = {};

    files.forEach(file => {
        const relativeDir = path.relative(rootDir, file.dir);
        const dirKey = relativeDir || '.';

        if (!dirSizes[dirKey]) {
            dirSizes[dirKey] = {
                totalSize: 0,
                totalSizeKB: '0.00',
                displayPath: relativeDir || rootDir,
                fullPath: file.dir
            };
        }

        dirSizes[dirKey].totalSize += file.size;
        dirSizes[dirKey].totalSizeKB = formatSize(dirSizes[dirKey].totalSize);
    });

    return dirSizes;
}


/**
 * 构建树状结构并计算层级占比
 */
function buildTreeStructure(dirSizes, rootDir) {
    const rootName = path.basename(rootDir);
    const tree = {
        name: rootName,
        fullPath: rootDir,
        relativePath: rootName,
        sizeKB: dirSizes['.']?.totalSizeKB || '0.00',
        children: [],
        level: 0
    };

    // 构建基础树结构
    const subDirs = Object.entries(dirSizes).filter(([key]) => key !== '.');
    subDirs.forEach(([relPath, dirInfo]) => {
        const pathParts = relPath.split(path.sep);
        let currentLevel = tree.children;
        let currentRelativePath = rootName;

        pathParts.forEach((part, index) => {
            currentRelativePath = path.join(currentRelativePath, part);
            let existingNode = currentLevel.find(node => node.name === part);
            
            if (!existingNode) {
                const isLastPart = index === pathParts.length - 1;
                existingNode = {
                    name: part,
                    fullPath: path.join(rootDir, relPath),
                    relativePath: currentRelativePath,
                    sizeKB: isLastPart ? dirInfo.totalSizeKB : '0.00',
                    children: [],
                    level: index + 1
                };
                currentLevel.push(existingNode);
            }
            currentLevel = existingNode.children;
        });
    });

    // 计算中间目录体积
    function calculateIntermediateSizes(node) {
        if (node.children.length === 0) return parseFloat(node.sizeKB);
        
        let total = parseFloat(node.sizeKB);
        node.children.forEach(child => {
            total += calculateIntermediateSizes(child);
        });
        
        node.sizeKB = total.toFixed(2);
        return total;
    }
    calculateIntermediateSizes(tree);

    // 计算层级占比
    function calculateLevelPercentage(node, parentSize) {
        if (node.level === 0) {
            node.percentage = '100.0';
        } else {
            const nodeSize = parseFloat(node.sizeKB) || 0;
            node.percentage = parentSize > 0 
                ? formatPercent((nodeSize / parentSize) * 100) 
                : '0.0';
        }

        node.children.forEach(child => {
            calculateLevelPercentage(child, parseFloat(node.sizeKB) || 0);
        });
    }
    calculateLevelPercentage(tree, 0);

    return tree;
}


/**
 * 打印树状体积图
 */
function printTreeVisualization(tree) {
    console.log(`\x1B[32m${tree.name} [${tree.sizeKB} KB]\x1B[0m`);
    
    function printChildren(children, parentPrefix = '') {
        children.forEach((child, index) => {
            const isLast = index === children.length - 1;
            const linePrefix = isLast ? '└─' : '├─';
            const childPrefix = isLast ? '  ' : '│ ';
            
            console.log(`${parentPrefix}${linePrefix} \x1B[36m${child.name} [${child.sizeKB} KB]\x1B[0m`);
            
            if (child.children.length > 0) {
                printChildren(child.children, parentPrefix + childPrefix);
            }
        });
    }

    printChildren(tree.children, '  ');
}


/**
 * 收集所有节点，按层级排序
 */
function collectNodesByLevel(tree) {
    const nodesByLevel = {};
    
    function traverse(node) {
        if (!nodesByLevel[node.level]) {
            nodesByLevel[node.level] = [];
        }
        
        if (node.level === 0 || parseFloat(node.sizeKB) > 0 || node.children.length > 0) {
            nodesByLevel[node.level].push({
                relativePath: node.relativePath,
                fullPath: node.fullPath,
                percentage: node.percentage
            });
        }
        
        node.children.forEach(child => traverse(child));
    }
    
    traverse(tree);
    return nodesByLevel;
}


/**
 * 打印层级百分比
 */
function printLevelPercentages(nodesByLevel) {
    const sortedLevels = Object.keys(nodesByLevel).sort((a, b) => parseInt(a) - parseInt(b));
    
    sortedLevels.forEach(level => {
        console.log(`\n\x1B[34m===== 第 ${level} 层级目录 =====\x1B[0m`);
        nodesByLevel[level].forEach(node => {
            console.log(`${node.relativePath} → ${node.percentage}%`);
        });
    });
}


/**
 * 主函数
 */
async function main() {
    try {
        let targetDir = process.cwd();
        if (process.argv.length > 2) {
            targetDir = path.resolve(process.argv[2]);
        }

        console.log(`\n📂 开始扫描目录：${targetDir}`);
        const files = await scanDirectory(targetDir);
        
        if (files.length === 0) {
            console.log(`\x1B[33mℹ️  未发现目标文件（支持类型：${targetExtensions.join(', ')}\x1B[0m`);
            return;
        }

        const dirSizes = calculateDirSizes(files, targetDir);
        const tree = buildTreeStructure(dirSizes, targetDir);
        const nodesByLevel = collectNodesByLevel(tree);

        // 先输出树状体积图
        console.log('\n📊 目录体积树状图：');
        printTreeVisualization(tree);
        
        // 再输出层级百分比
        console.log('\n📊 目录层级占比（每行一个层级）：');
        printLevelPercentages(nodesByLevel);
        
        console.log('\n✅ 扫描完成！');

    } catch (err) {
        console.error(`\x1B[31m❌ 执行失败：${err.message}\x1B[0m`);
        process.exit(1);
    }
}


// 启动程序
main();
